from cassandra.cluster import Cluster
from cassandra import ConsistencyLevel
from pylab import *
from matplotlib.dates import AutoDateLocator, DateFormatter

if __name__ == '__main__':
    cluster = Cluster(protocol_version = 4)
    session = cluster.connect()
    session.set_keyspace('wmcloud')
    # query data from wmcloud
    stock_info_stmt = session.prepare('select * from wmcloud.sz_sh_stock_daily_trade_info where ticker = ? order by tradedate desc limit ?')
    stock_info_stmt.consistency_level = ConsistencyLevel.QUORUM

    stock_ticker = '600000'
    stock_info_row_list = session.execute(stock_info_stmt, [stock_ticker, 200])
    close_price_list = []
    highest_price_list = []
    lowest_price_list = []
    open_price_list = []
    turn_over_volume_list = []
    tradedate_list = []
    for stock_info_row in stock_info_row_list:
        open_price_list.append(stock_info_row.openprice)
        close_price_list.append(stock_info_row.closeprice)
        highest_price_list.append(stock_info_row.highestprice)
        lowest_price_list.append(stock_info_row.lowestprice)
        turn_over_volume_list.append(stock_info_row.turnovervol)
        tradedate_list.append(stock_info_row.tradedate)

    session.shutdown()
    cluster.shutdown()

    autodates = AutoDateLocator()
    timefmt = DateFormatter('%Y-%m-%d')

    close_price_ax = subplot(211)
    close_price_ax.set_title('Stock Price Trend %s' % (stock_ticker))
    close_price_ax.set_xlabel('Time')
    close_price_ax.set_ylabel('Price')
    close_price_ax.xaxis.set_major_locator(autodates)
    close_price_ax.xaxis.set_major_formatter(timefmt)
    close_price_ax.plot(date2num(tradedate_list), close_price_list, linestyle = '-', label = '%s close price' % (stock_ticker), color = 'y')
    close_price_ax.plot(date2num(tradedate_list), highest_price_list, linestyle = '-', label = '%s highest price' % (stock_ticker), color = 'r')
    close_price_ax.plot(date2num(tradedate_list), lowest_price_list, linestyle = '-', label = '%s lowest price' % (stock_ticker), color = 'b')
    close_price_ax.plot(date2num(tradedate_list), open_price_list, linestyle = '-', label = '%s openprice' % (stock_ticker), color = 'c')

    turn_over_vol_ax = subplot(212)
    turn_over_vol_ax.set_title('Stock Turnovervol Trend %s' % (stock_ticker))
    turn_over_vol_ax.set_xlabel('Time')
    turn_over_vol_ax.set_ylabel('Turnovervol')
    turn_over_vol_ax.xaxis.set_major_locator(autodates)
    turn_over_vol_ax.xaxis.set_major_formatter(timefmt)
    # turn_over_vol_ax.plot(date2num(tradedate_list), turn_over_volume_list, linestyle = '-')
    for inx, tradedate in enumerate(tradedate_list):
        if open_price_list[inx] > close_price_list[inx]:
            turn_over_vol_ax.bar(date2num(tradedate), turn_over_volume_list[inx], facecolor = '#99ff99', edgecolor = 'white')
        elif open_price_list[inx] < close_price_list[inx]:
            turn_over_vol_ax.bar(date2num(tradedate), turn_over_volume_list[inx], facecolor = '#ff9999', edgecolor = 'white')
        else:
            turn_over_vol_ax.bar(date2num(tradedate), turn_over_volume_list[inx], facecolor = '#9999ff', edgecolor = 'white')

    #股价和成交量是输入也是导致未来的股价和成交量的因素
    #在股市中应该存在一个正反馈效应，通过其它变量的影响
    #最终对这个市场进行调节
    #除了正反馈，还有一个独立输入源就是市场参与者，也就是
    #人心的评估要作为一个决定性输入变量参与进来，也就是当前
    #的正反馈参数和一个评估参数的综合结果决定未来的走势
    #如果越准确，那么对于未来的走势判断就会越准确．

    fig = gcf()
    fig.autofmt_xdate()
    show()